from django.shortcuts import render
from django.http import FileResponse, Http404
from django.conf import settings
import os
# Create your views here.
from rest_framework import viewsets
from .models import Task, KnowledgeBase, KnowledgeBaseFile
from .serializers import TaskSerializer, KnowledgeBaseSerializer, KnowledgeBaseFileSerializer
from rest_framework.views import APIView
from rest_framework.response import Response
from rest_framework import status
from sentence_transformers import SentenceTransformer
from pymilvus import connections, Collection, FieldSchema, CollectionSchema, DataType, utility
import numpy as np
from .serializers import TaskSerializer, TextEmbeddingSerializer, QuestionSerializer
from langchain_openai import ChatOpenAI
from langchain_community.vectorstores import Milvus
from langchain.prompts import PromptTemplate
from langchain_core.output_parsers import StrOutputParser
from langchain_community.graphs import Neo4jGraph
from text2vec import SentenceModel
from langchain.chains import GraphCypherQAChain
from PyPDF2 import PdfReader
import io
from rest_framework.decorators import action
from rest_framework import viewsets
from .models import ProcessingData
from .serializers import ProcessingDataSerializer  # 确保你有相应的序列化器
from drf_yasg.utils import swagger_auto_schema
from drf_yasg import openapi
from rest_framework.parsers import MultiPartParser, FormParser
# 任务测试
class TaskViewSet(viewsets.ModelViewSet):
    queryset = Task.objects.all()
    serializer_class = TaskSerializer

# 文本向量化
class TextEmbeddingView(APIView):
    def __init__(self):
        super().__init__()
        # 连接Milvus
        connections.connect(
            alias="default", 
            host="218.199.129.36",
            port="19530"
        )
        # 加载模型
        self.model = SentenceTransformer(r"D:\400code\python\llm_maxkb\shibing624_text2vec-base-chinese\shibing624_text2vec-base-chinese")
        
        # 确保collection存在
        self.create_collection_if_not_exists()
    
    def create_collection_if_not_exists(self):
        collection_name = "text_embeddings_file"
        
        if not utility.has_collection(collection_name):
            # 定义字段
            fields = [
                FieldSchema(name="id", dtype=DataType.INT64, is_primary=True, auto_id=True),
                FieldSchema(name="text", dtype=DataType.VARCHAR, max_length=5000),
                FieldSchema(name="embedding", dtype=DataType.FLOAT_VECTOR, dim=768)  # text2vec-base-chinese 输出维度为768
            ]
            
            # 创建collection
            schema = CollectionSchema(fields)
            self.collection = Collection(name=collection_name, schema=schema)
            
            # 创建索引
            index_params = {
                "metric_type": "L2",
                "index_type": "IVF_FLAT",
                "params": {"nlist": 1024}
            }
            self.collection.create_index(field_name="embedding", index_params=index_params)
        else:
            self.collection = Collection(name=collection_name)
    
    def post(self, request):
        if 'pdf_file' not in request.FILES:
            return Response({"error": "未提供PDF文件"}, status=status.HTTP_400_BAD_REQUEST)
        
        pdf_file = request.FILES['pdf_file']
        pdf_reader = PdfReader(io.BytesIO(pdf_file.read()))
        
        full_text = ""
        for page in pdf_reader.pages:
            full_text += page.extract_text()
        
        # 按4096字符一组分割文本
        text_chunks = [full_text[i:i+4096] for i in range(0, len(full_text), 4096)]
        
        all_texts = []
        all_embeddings = []
        
        for chunk in text_chunks:
            # 生成文本嵌入
            embedding = self.model.encode(chunk)
            
            all_texts.append(chunk)
            all_embeddings.append(embedding.tolist())
        
        # 准备插入数据
        data = [
            all_texts,  # text field
            all_embeddings  # embedding field
        ]
        
        # 插入数据到Milvus
        self.collection.insert(data)
        self.collection.flush()
        
        return Response({
            "message": "文本嵌入已成功存储",
            "num_chunks": len(text_chunks)
        }, status=status.HTTP_201_CREATED)

class QAView(APIView):
    def __init__(self):
        super().__init__()
        # 环境变量设置
        os.environ["OPENAI_API_KEY"] = "sk-4jPlYAHFMLAAgWPWHKjFcAcRLHXBrhGCEKPTwjhRd5XVvX2Y"
        os.environ["NEO4J_URI"] = "bolt://218.199.129.36:7687"
        os.environ["NEO4J_USERNAME"] = "neo4j"
        os.environ["NEO4J_PASSWORD"] = "12345678"
        
        # 初始化LLM
        self.llm = ChatOpenAI(
            temperature=0,
            model_name="gpt-4",
            base_url="https://api.chatanywhere.tech/v1"
        )
        
        # 初始化Neo4j
        self.graph = Neo4jGraph(
            url=os.getenv("NEO4J_URI"),
            username=os.getenv("NEO4J_USERNAME"),
            password=os.getenv("NEO4J_PASSWORD")
        )
        
        # 初始化向量存储
        self.init_vector_store()
        
        # 初始化提示模板
        self.init_prompts()
        
        # 初始化链
        self.init_chains()
        
    def init_vector_store(self):
        class SentenceEmbedding:
            def __init__(self, model):
                self.model = model
            
            def embed_documents(self, texts):
                return self.model.encode(texts)
            
            def embed_query(self, texts):
                return self.model.encode(texts)
        
        model = SentenceModel('shibing624/text2vec-base-chinese')
        embedding_function = SentenceEmbedding(model)
        
        # self.vectorstore = Milvus(
        #     embedding_function=embedding_function,
        #     connection_args={"host": "218.199.129.36", "port": "19530"},
        #     collection_name="rag_milvus2",
        # )
        # self.retriever = self.vectorstore.as_retriever()
        try:
            self.vectorstore = Milvus(
                embedding_function=embedding_function,
                connection_args={
                    "uri": "http://218.199.129.36:19530",  # 使用 http 协议
                    "user": "",        # 如果有用户名
                    "password": "",    # 如果有密码
                    "secure": False    # 如果不需要 SSL
                },
                collection_name="rag_milvus2",
            )
            print("Successfully connected to Milvus")
        except Exception as e:
            print(f"Failed to connect to Milvus: {str(e)}")
            raise
            
        self.retriever = self.vectorstore.as_retriever()
    
    def init_prompts(self):
        self.cypher_prompt = PromptTemplate(
            template="""You are an expert at generating Cypher queries for Neo4j.
            Use the following schema to generate a Cypher query that answers the given question.
            Make the query flexible by using case-insensitive matching and partial string matching where appropriate.
            Focus on searching paper titles as they contain the most relevant information.

            Schema:
            {schema}

            Question: {question}

            Cypher Query:""",
            input_variables=["schema", "question"],
        )
        
        self.qa_prompt = PromptTemplate(
            template="""You are an assistant for question-answering tasks. 
            Use the following Cypher query results to answer the question. If you don't know the answer, just say that you don't know. 
            Use three sentences maximum and keep the answer concise. If topic information is not available, focus on the paper titles.

            Question: {question} 
            Cypher Query: {query}
            Query Results: {context} 

            Answer:""",
            input_variables=["question", "query", "context"],
        )
        
        self.prompt_vector = PromptTemplate(
            template="""You are an assistant for question-answering tasks. 
            Use the following pieces of retrieved context to answer the question. If you don't know the answer, just say that you don't know. 
            Use three sentences maximum and keep the answer concise:
            Question: {question} 
            Context: {context} 
            Answer:""",
            input_variables=["question", "context"],
        )
        
        self.final_prompt = PromptTemplate(
            template="""You are an assistant for question-answering tasks. 
            Use the following pieces of retrieved context from a vector store and a graph database to answer the question. If you don't know the answer, just say that you don't know. 
            Use three sentences maximum and keep the answer concise:
            Question: {question} 
            Vector Context: {context} 
            Graph Context: {graph_context}
            Answer:""",
            input_variables=["question", "context", "graph_context"],
        )
    
    def init_chains(self):
        self.rag_chain = self.prompt_vector | self.llm | StrOutputParser()
        self.graph_rag_chain = GraphCypherQAChain.from_llm(
            cypher_llm=self.llm,
            qa_llm=self.llm,
            validate_cypher=True,
            graph=self.graph,
            verbose=True,
            return_intermediate_steps=True,
            return_direct=True,
            cypher_prompt=self.cypher_prompt,
            qa_prompt=self.qa_prompt,
            allow_dangerous_requests=True,
        )
        self.composite_chain = self.final_prompt | self.llm | StrOutputParser()
    
    def post(self, request):
        serializer = QuestionSerializer(data=request.data)
        if serializer.is_valid():
            question = serializer.validated_data['question']
            
            try:
                # 获取向量数据库的结果
                docs_qu = self.retriever.invoke(question)
                vector_context = self.rag_chain.invoke({"context": docs_qu, "question": question})
                
                # 获取图数据库的结果
                graph_context = self.graph_rag_chain.invoke({"query": question})
                
                # 综合结果
                answer = self.composite_chain.invoke({
                    "question": question,
                    "context": vector_context,
                    "graph_context": graph_context
                })
                
                return Response({
                    "question": question,
                    "answer": answer
                }, status=status.HTTP_200_OK)
                
            except Exception as e:
                return Response({
                    "error": str(e)
                }, status=status.HTTP_500_INTERNAL_SERVER_ERROR)
                
        return Response(serializer.errors, status=status.HTTP_400_BAD_REQUEST)

# 知识库 资料库

class KnowledgeBaseViewSet(viewsets.ModelViewSet):
    queryset = KnowledgeBase.objects.all()
    serializer_class = KnowledgeBaseSerializer
    # 设置解析器类
    parser_classes = [MultiPartParser, FormParser]
    def list(self, request, *args, **kwargs):
        # 在这里可以添加自定义的查询逻辑
        queryset = self.filter_queryset(self.get_queryset())

        # 可以根据需要对查询集进行进一步处理

        serializer = self.get_serializer(queryset, many=True)
        # 可以对序列化后的数据进行处理

        return Response({'code': 20000, 'data': serializer.data})
    def create(self, request):
        serializer = self.get_serializer(data=request.data)
        if serializer.is_valid():
            knowledge_base = serializer.save()  # 保存知识库
            knowledge_base.directory = f"knowledge_base_{knowledge_base.id}"  # 生成目录名称
            knowledge_base.save()  # 更新目录
            return Response({"code": 20000, "data": serializer.data}, status=status.HTTP_201_CREATED)
        return Response({"code": status.HTTP_400_BAD_REQUEST, "data": serializer.errors}, status=status.HTTP_400_BAD_REQUEST)

    def update(self, request, pk=None):
        knowledge_base = self.get_object()
        serializer = self.get_serializer(knowledge_base, data=request.data, partial=True)
        if serializer.is_valid():
            serializer.save()
            return Response({"code": 20000, "data": serializer.data}, status=status.HTTP_200_OK)
        return Response({"code": status.HTTP_400_BAD_REQUEST, "data": serializer.errors}, status=status.HTTP_400_BAD_REQUEST)

    def destroy(self, request, pk=None):
        knowledge_base = self.get_object()
        knowledge_base.delete()
        return Response({"code": 20000, "data": {"message": "文件夹已成功删除"}}, status=status.HTTP_204_NO_CONTENT)

    def upload_file(self, request, pk=None):
        knowledge_base = self.get_object()
        files = request.FILES.getlist('files')  # 获取多个文件
        uploaded_files = []

        for file in files:
            knowledge_base_file = KnowledgeBaseFile(knowledge_base=knowledge_base, file=file)
            knowledge_base_file.save()
            uploaded_files.append({
                "id": knowledge_base_file.id,  # 返回文件的 ID
                "file_name": knowledge_base_file.file.name,
                "file_url": request.build_absolute_uri(knowledge_base_file.file.url)
            })

        return Response({"code": 20000, "data": {"uploaded_files": uploaded_files}}, status=status.HTTP_201_CREATED)

    @action(detail=True, methods=['get'], url_path='download-file/(?P<file_id>[^/.]+)')
    def download_file(self, request, pk=None, file_id=None):
        knowledge_base = self.get_object()
        try:
            # 获取指定的文件
            knowledge_base_file = KnowledgeBaseFile.objects.get(id=file_id, knowledge_base=knowledge_base)
            file_path = os.path.join(settings.MEDIA_ROOT, knowledge_base_file.file.name)
            return FileResponse(open(file_path, 'rb'), as_attachment=True)
        except KnowledgeBaseFile.DoesNotExist:
            return Response({"code": status.HTTP_404_NOT_FOUND, "data": {"error": "文件不存在"}}, status=status.HTTP_404_NOT_FOUND)
        except Exception as e:
            return Response({"code": status.HTTP_500_INTERNAL_SERVER_ERROR, "data": {"error": str(e)}}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)


    def delete_file(self, request, pk=None):
        knowledge_base = self.get_object()
        if knowledge_base.file:
            os.remove(os.path.join(settings.MEDIA_ROOT, knowledge_base.file.name))  # 删除文件
            knowledge_base.file = None  # 清空文件字段
            knowledge_base.save()
            return Response({"code": 20000, "data": {"message": "文件已成功删除"}}, status=status.HTTP_204_NO_CONTENT)
        return Response({"code": status.HTTP_404_NOT_FOUND, "data": {"error": "文件不存在"}}, status=status.HTTP_404_NOT_FOUND)

    
    @swagger_auto_schema(
        operation_description="批量上传文件",  # 接口描述
        manual_parameters=[
            openapi.Parameter(
                'files',  # 参数名称
                openapi.IN_FORM,  # 参数位置
                description="上传的文件列表",  # 参数描述
                type=openapi.TYPE_FILE,  # 参数类型
                required=True,  # 是否必填
            ),
            openapi.Parameter('resource_code', openapi.IN_FORM, description="资源编号", type=openapi.TYPE_STRING),
            openapi.Parameter('resource_name', openapi.IN_FORM, description="资源名称", type=openapi.TYPE_STRING),
            openapi.Parameter('resource_type', openapi.IN_FORM, description="资源类型", type=openapi.TYPE_STRING),
            openapi.Parameter('resource_source', openapi.IN_FORM, description="资源来源", type=openapi.TYPE_STRING),
            openapi.Parameter('resource_quality', openapi.IN_FORM, description="资源质量", type=openapi.TYPE_STRING),
            openapi.Parameter('version_number', openapi.IN_FORM, description="版本号", type=openapi.TYPE_STRING),
            openapi.Parameter('resource_description', openapi.IN_FORM, description="资源描述", type=openapi.TYPE_STRING),
            openapi.Parameter('resource_contact_person', openapi.IN_FORM, description="联系人", type=openapi.TYPE_STRING),
            openapi.Parameter('contact_phone', openapi.IN_FORM, description="联系电话", type=openapi.TYPE_STRING),
            openapi.Parameter('contact_email', openapi.IN_FORM, description="联系邮箱", type=openapi.TYPE_STRING),
            openapi.Parameter('notes', openapi.IN_FORM, description="备注", type=openapi.TYPE_STRING),
        ]
    )
    @action(detail=True, methods=['post'], url_path='upload-files')
    def upload_files(self, request, pk=None):
        knowledge_base = self.get_object()
        files = request.FILES.getlist('files')  # 获取多个文件
        uploaded_files = []

        for file in files:
            knowledge_base_file = KnowledgeBaseFile(
                knowledge_base=knowledge_base,
                file=file,
                resource_code=request.data.get('resource_code', ''),
                resource_name=request.data.get('resource_name', ''),
                resource_type=request.data.get('resource_type', ''),
                resource_source=request.data.get('resource_source', ''),
                resource_quality=request.data.get('resource_quality', ''),
                version_number=request.data.get('version_number', ''),
                resource_description=request.data.get('resource_description', ''),
                resource_contact_person=request.data.get('resource_contact_person', ''),
                contact_phone=request.data.get('contact_phone', ''),
                contact_email=request.data.get('contact_email', ''),
                notes=request.data.get('notes', '')
            )
            knowledge_base_file.save()
            uploaded_files.append({
                "file_name": knowledge_base_file.file.name,
                "file_url": request.build_absolute_uri(knowledge_base_file.file.url)
            })

        return Response({"code": 20000, "data": {"uploaded_files": uploaded_files}}, status=status.HTTP_201_CREATED)

    @action(detail=True, methods=['delete'], url_path='delete-files')
    def delete_files(self, request, pk=None):
        knowledge_base = self.get_object()
        file_ids = request.data.get('file_ids', [])

        if not file_ids:
            return Response({"error": "未提供文件ID"}, status=status.HTTP_400_BAD_REQUEST)

        deleted_files = []
        for file_id in file_ids:
            try:
                knowledge_base_file = KnowledgeBaseFile.objects.get(id=file_id, knowledge_base=knowledge_base)
                knowledge_base_file.delete()
                deleted_files.append(file_id)
            except KnowledgeBaseFile.DoesNotExist:
                return Response({"code": status.HTTP_404_NOT_FOUND, "data": {"error": f"文件ID {file_id} 不存在"}}, status=status.HTTP_404_NOT_FOUND)

        return Response({"code": 20000, "data": {"deleted_files": deleted_files}}, status=status.HTTP_204_NO_CONTENT)


    @action(detail=True, methods=['get'], url_path='files')
    def list_files(self, request, pk=None):
        knowledge_base = self.get_object()  # 获取当前知识库对象
        files = knowledge_base.knowledgebasefile_set.all()  # 获取与知识库相关的所有文件
        file_info = []

        for file in files:
            file_info.append({
                "id": file.id,
                "file_name": file.file.name,
                "file_url": request.build_absolute_uri(file.file.url),  # 生成文件的完整 URL
                "resource_code": file.resource_code,  # 资料编号
                "resource_name": file.resource_name,  # 资料名称
                "resource_type": file.resource_type,  # 资料分类
                "resource_source": file.resource_source,  # 资料来源
                "resource_quality": file.resource_quality,  # 资料质量
                "version_number": file.version_number,  # 版本号
                "resource_description": file.resource_description,  # 资料描述
                "resource_contact_person": file.resource_contact_person,  # 资料联系人
                "contact_phone": file.contact_phone,  # 联系电话
                "contact_email": file.contact_email,  # 联系Email
                "notes": file.notes,  # 备注
            })

        return Response({"code": 20000, "data": {"files": file_info}}, status=status.HTTP_200_OK)


class ProcessingDataViewSet(viewsets.ModelViewSet):
    queryset = ProcessingData.objects.all()
    serializer_class = ProcessingDataSerializer


class ProcessingDataViewSet(viewsets.ModelViewSet):
    queryset = ProcessingData.objects.all()
    serializer_class = ProcessingDataSerializer
    
    @action(detail=False, methods=['post'], url_path='upload-files')
    def upload_files(self, request, pk=None):
        try:
            # 通过 pk 手动查询 KnowledgeBase 对象
            knowledge_base = KnowledgeBase.objects.get(pk=pk)
        except KnowledgeBase.DoesNotExist:
            return Response({"error": "KnowledgeBase with the given ID does not exist."}, status=status.HTTP_404_NOT_FOUND)
        files = request.FILES.getlist('files')  # 获取多个文件
        uploaded_files = []

        for file in files:
            processing_data = ProcessingData(
                file=file,
                result_code=request.data.get('result_code', ''),
                result_name=request.data.get('result_name', ''),
                result_type=request.data.get('result_type', ''),
                method_used=request.data.get('method_used', ''),
                method_parameters=request.data.get('method_parameters', ''),
                instructions=request.data.get('instructions', ''),
                notes=request.data.get('notes', ''),
                original_file_id=request.data.get('original_file_id', ''),
                resource_directory=request.data.get('resource_directory', ''),
                # knowledge_base_id=request.data.get('knowledge_base_id', None)  # 处理知识库ID
                knowledge_base=knowledge_base  # 直接关联 KnowledgeBase 对象
            )
            processing_data.save()
            uploaded_files.append({
                "file_name": processing_data.file.name,
                "file_url": request.build_absolute_uri(processing_data.file.url)
            })
        return Response({"code": 20000, "data": {"uploaded_files": uploaded_files}}, status=status.HTTP_201_CREATED)
    
    @action(detail=True, methods=['get'], url_path='files')
    def list_files(self, request, pk=None):
        try:
            # 通过 pk 手动查询 KnowledgeBase 对象
            knowledge_base = KnowledgeBase.objects.get(pk=pk)
        except KnowledgeBase.DoesNotExist:
            return Response({"error": "KnowledgeBase with the given ID does not exist."}, status=status.HTTP_404_NOT_FOUND)
        files = knowledge_base.processing_data_files.all()  # 获取与知识库相关的所有文件
        file_info = []

        for file in files:
            file_info.append({
                "id": file.id,
                "file_url": request.build_absolute_uri(file.file.url),  # 生成文件的完整 URL
                "result_code": file.result_code,  # 资料编号
                "result_name": file.result_name,  # 资料名称
                "result_type": file.result_type,  # 资料分类
                "method_used": file.method_used,  # 资料来源
                "method_parameters": file.method_parameters,  # 资料质量
                "instructions": file.instructions,  # 版本号
                "resource_directory": file.resource_directory,  # 联系电话
                "original_file_id": file.original_file_id,  # 联系Email
                "notes": file.notes,  # 备注
            })

        return Response({"code": 20000, "data": {"files": file_info}}, status=status.HTTP_200_OK)
    @action(detail=True, methods=['get'], url_path='download-file/(?P<file_id>[^/.]+)')
    def download_file(self, request, pk=None, file_id=None):
        try:
            # 通过 pk 手动查询 KnowledgeBase 对象
            knowledge_base = KnowledgeBase.objects.get(pk=pk)
        except KnowledgeBase.DoesNotExist:
            return Response({"error": "KnowledgeBase with the given ID does not exist."}, status=status.HTTP_404_NOT_FOUND)
        try:
            # 获取指定的文件
            processingData = ProcessingData.objects.get(id=file_id, knowledge_base=knowledge_base)
            file_path = os.path.join(settings.MEDIA_ROOT, processingData.file.name)
            return FileResponse(open(file_path, 'rb'), as_attachment=True)
        except ProcessingData.DoesNotExist:
            return Response({"code": status.HTTP_404_NOT_FOUND, "data": {"error": "文件不存在"}}, status=status.HTTP_404_NOT_FOUND)
        except Exception as e:
            return Response({"code": status.HTTP_500_INTERNAL_SERVER_ERROR, "data": {"error": str(e)}}, status=status.HTTP_500_INTERNAL_SERVER_ERROR)